package io.gitee.zhangbinhub.acp.boot.websocket.reactive

import io.gitee.zhangbinhub.acp.boot.log.LogAdapter
import org.springframework.web.reactive.socket.WebSocketHandler
import org.springframework.web.reactive.socket.WebSocketMessage
import org.springframework.web.reactive.socket.WebSocketSession
import reactor.core.publisher.FluxSink
import reactor.core.publisher.Mono
import java.nio.ByteBuffer

abstract class AcpWebSocketHandlerReactive(@JvmField protected val logAdapter: LogAdapter) : WebSocketHandler {
    @JvmField
    protected var session: WebSocketSession? = null

    @JvmField
    protected var outFluxSink: FluxSink<WebSocketMessage>? = null
    private fun doSendMessage(message: WebSocketMessage) = outFluxSink!!.next(message)

    fun sendMessage(message: ByteBuffer) {
        if (beforeSendMessage()) {
            doSendMessage(session!!.binaryMessage { factory -> factory.wrap(message) })
        }
    }

    fun sendMessage(message: String) {
        if (beforeSendMessage()) {
            doSendMessage(session!!.textMessage(message))
        }
    }

    private fun beforeSendMessage(): Boolean =
        if (session != null && session!!.isOpen) {
            true
        } else {
            logAdapter.error("WebSocket session is null or closed!")
            false
        }

    protected fun doReceiveMessage(session: WebSocketSession): Mono<Void> =
        session.receive().mapNotNull { message ->
            message
        }.doOnNext { message ->
            when (message.type) {
                WebSocketMessage.Type.BINARY -> {
                    ByteArray(message.payload.readableByteCount()).apply {
                        message.payload.read(this)
                    }.apply {
                        receiveMessage(session, ByteBuffer.wrap(this))
                    }
                }

                WebSocketMessage.Type.TEXT -> {
                    receiveMessage(session, message.payloadAsText)
                }

                WebSocketMessage.Type.PING -> {
                    outFluxSink!!.next(session.pongMessage { factory -> factory.wrap(ByteArray(0)) })
                }

                WebSocketMessage.Type.PONG -> {
                    logAdapter.debug("WebSocket server receive pong message")
                }

                else -> throw IllegalStateException("Unexpected WebSocket message type: ${message.type}")
            }
        }.doOnError { throwable ->
            sessionError(session, throwable)
        }.then()

    abstract fun sessionOpen(session: WebSocketSession)
    abstract fun sessionClose(session: WebSocketSession)
    abstract fun sessionError(session: WebSocketSession?, throwable: Throwable?)
    abstract fun receiveMessage(session: WebSocketSession, message: ByteBuffer)
    abstract fun receiveMessage(session: WebSocketSession, message: String)
}